import torch
from torch import nn
import torch.nn.functional as F

class Attention(nn.Module):
    def __init__(self,query_size, key_size, value_size1, value_size2, output_size):
        super(Attention, self).__init__()
        self.query_size = query_size
        self.key_size = key_size
        self.value_size1 = value_size1
        self.value_size2 = value_size2
        self.output_size = output_size
        # 第一种方式
        # 先拼接 然后进行线性变换 然后softmax
        # 最后乘V

        # 初始化注意力机制第一步
        # 两个size相加，是直接把矩阵拼接
        # 拼接后进行线性变换使用 （64,32）
        self.attn = nn.Linear(self.query_size + self.key_size, self.value_size1)

        # 最后乘V后 输出使用
        # 初始化注意力机制第三步
        # 线性变换 （96,64）
        self.attn_combine = nn.Linear(self.query_size + self.value_size2, self.output_size)

    #     Q K 都是三维数据 维度是相同的 （1,1,32）
    def forward(self, Q, K, V):
        # 采用上述第一种计算规则

        # 先进性QK的拼接以及线性变换，再经过softmax处理获得结果
        # 这里QKV都是三维张量
        # （1,32） （1,32） 在维度1 上的cat的维度为（1,64）
        # （1,64） * （64,32）
        # 结果为（1,32） 这是第一个线性变换的结果 然后在32这个维度上进行 softmax 最后的维度还是不变的
        attn_weights = F.softmax(self.attn(torch.cat((Q[0], K[0]), 1)), dim=1)

        # 然后将结果 与 V相乘 （1,1,32） @ （1,32,64）= （1,1,64）
        attn_applied = torch.bmm(attn_weights.unsqueeze(0), V)

        # 第二步，将Q与第一步的结果再进行拼接
        # （1,32） （1，64） 在第1个维度进行拼接 结果为（1,96）
        output = torch.cat((Q[0], attn_applied[0]), 1)

        # （1,96） * （96,64）=（1,64）
        # 经过unsqueeze (1,1,64)
        # 第三步，得到输出
        output = self.attn_combine(output).unsqueeze(0)
        return output, attn_weights
        # (1, 1, 64) (1,32)

query_size = 32
key_size = 32
value_size1 = 32
value_size2 = 64
output_size = 64
attn = Attention(query_size, key_size, value_size1, value_size2, output_size)
Q = torch.randn(1, 1, query_size)
print("---")
K = torch.randn(1, 1, key_size)
s = torch.cat((Q[0], K[0]), 1)
V = torch.randn(1, value_size1, value_size2)
out = attn(Q, K, V)
